#!/usr/bin/env python2
# -*- coding: UTF-8 -*-
import sys
import os

sys.path.append(os.path.abspath(os.getcwd()))

import requests
import threading
from common.helper import normal_laf

recursive_lock = threading.Lock()
write_result_lock = threading.Lock()

defaultCheckers = (
    lambda r: r.status_code == 200,
)


def visit(url, params=None, headers=None, checkers=defaultCheckers):
    try:
        r = requests.get(url, params=params, headers=headers)
        for checker in checkers:
            if not checker(r):
                return False
        return True
    except Exception as e:
        return False
    pass


def travel():
    while True:
        path = get_path()
        if path is None:
            break
        url = site + path
        print(threading.current_thread().name + ' : ' + url)
        if visit(url, module.params, module.headers,
                 module.checkers if module.checkers is not None else defaultCheckers):
            write_result_lock.acquire()
            resultFile.write(url + '\n')
            write_result_lock.release()


def get_path():
    recursive_lock.acquire()
    if len(paths) > 0:
        path = paths.pop()
    else:
        path = None
    recursive_lock.release()
    return path


paths = []
module = None
site = None


def main():
    global module, paths, site

    module = __import__(raw_input('请输入前置脚本: '))
    site = raw_input('请输入目标网站: ').strip()
    thread_count = raw_input('请输入开启的线程数: ')
    dict_file = raw_input('请输入目录字典文件: ')
    print('目标站点:%s ; 开启线程数: %s ; 目录字典: %s' % (site, thread_count, dict_file))

    paths = normal_laf(os.path.join(os.path.abspath(os.path.dirname(__file__)), 'dict', dict_file))

    threads = []
    for i in range(1, int(thread_count)):
        t = threading.Thread(target=travel)
        t.start()
        threads.append(t)

    for t in threads:
        t.join()

    print('Done!')


resultFile = open(os.path.join(os.path.abspath(os.path.dirname(__file__)), 'result'), 'w+')

if __name__ == '__main__':
    main()
